import { Vector2 } from 'three';

export class PoissonDiskSampling {
    private grid: (Vector2 | null)[][];
    private activePoints: Vector2[];
    private cellSize: number;
    private radius: number;
    private width: number;
    private height: number;
    private numSamplesBeforeRejection: number;

    constructor(width: number, height: number, radius: number, numSamplesBeforeRejection = 30) {
        this.width = width;
        this.height = height;
        this.radius = radius;
        this.numSamplesBeforeRejection = numSamplesBeforeRejection;
        this.cellSize = radius / Math.sqrt(2);

        const cols = Math.ceil(width / this.cellSize);
        const rows = Math.ceil(height / this.cellSize);

        this.grid = Array(cols).fill(null).map(() => Array(rows).fill(null));
        this.activePoints = [];
    }

    private isValid(point: Vector2): boolean {
        if (point.x < 0 || point.x >= this.width || point.y < 0 || point.y >= this.height) {
            return false;
        }

        const cellX = Math.floor(point.x / this.cellSize);
        const cellY = Math.floor(point.y / this.cellSize);

        const searchStartX = Math.max(0, cellX - 2);
        const searchEndX = Math.min(this.grid.length - 1, cellX + 2);
        const searchStartY = Math.max(0, cellY - 2);
        const searchEndY = Math.min(this.grid[0].length - 1, cellY + 2);

        for (let x = searchStartX; x <= searchEndX; x++) {
            for (let y = searchStartY; y <= searchEndY; y++) {
                const point2 = this.grid[x][y];
                if (point2) {
                    const dist = Math.sqrt(
                        Math.pow(point.x - point2.x, 2) + 
                        Math.pow(point.y - point2.y, 2)
                    );
                    if (dist < this.radius) {
                        return false;
                    }
                }
            }
        }
        return true;
    }

    private addPoint(point: Vector2): void {
        const cellX = Math.floor(point.x / this.cellSize);
        const cellY = Math.floor(point.y / this.cellSize);
        this.grid[cellX][cellY] = point;
        this.activePoints.push(point);
    }

    public sample(numPoints: number): Vector2[] {
        // 初始点
        const firstPoint = new Vector2(
            Math.random() * this.width,
            Math.random() * this.height
        );
        
        this.addPoint(firstPoint);
        const points: Vector2[] = [firstPoint];

        while (points.length < numPoints && this.activePoints.length > 0) {
            const activeIndex = Math.floor(Math.random() * this.activePoints.length);
            const activePoint = this.activePoints[activeIndex];
            let found = false;

            for (let i = 0; i < this.numSamplesBeforeRejection; i++) {
                const angle = Math.random() * Math.PI * 2;
                const r = Math.random() * this.radius + this.radius;
                const newPoint = new Vector2(
                    activePoint.x + Math.cos(angle) * r,
                    activePoint.y + Math.sin(angle) * r
                );

                if (this.isValid(newPoint)) {
                    this.addPoint(newPoint);
                    points.push(newPoint);
                    found = true;
                    break;
                }
            }

            if (!found) {
                this.activePoints.splice(activeIndex, 1);
            }

            if (points.length >= numPoints) {
                break;
            }
        }

        return points;
    }
}
